#include "Master.h"
#include "mpi.h"
#include "Utils.h"

using namespace std;

// Job runner
/**************************************************************************************************************/
void Master :: run() {
    for(int i = 0; i < iterations_count; ++i) {
        initializeIteration();
        assignAllBlocksToSlaves();
        waitForEndOfComputations();
        finishIteration();
    }
    block_info_server.stop();
    sendTerminationSignal();
}
/**************************************************************************************************************/


// Iteration parts
/**************************************************************************************************************/
void Master :: initializeIteration() {
    todo.clear();
    in_progress.clear();
    for(int i = 0; i < blocks_count; ++i) {
        todo.insert(i);
    }
    
    IterationInfo :: swapContent(prev_info, current_info);
    current_info.clear();
}

void Master :: assignAllBlocksToSlaves() {
    int updated_block, idle_slave;
    
    while(todo.empty() == false) {
        recieveReadyToWorkSignal(updated_block, idle_slave);
        markBlockAsDone(updated_block, idle_slave);
        
        int next_block = chooseNextBlockFor(idle_slave);
        sendBlockAssignmentSignal(next_block, idle_slave);
        todo.erase(next_block);
        in_progress.insert(next_block);
    }
}

void Master :: waitForEndOfComputations() {
    int updated_block, idle_slave;
    
    while(in_progress.empty() == false) {
        recieveReadyToWorkSignal(updated_block, idle_slave);
        markBlockAsDone(updated_block, idle_slave);
    }
}

void Master :: finishIteration() {
    sendNextIterationSignal();
}
/**************************************************************************************************************/


// Block operators
/**************************************************************************************************************/
bool Master :: isValidBlock(int block_id) {
    return block_id >= 0 && block_id < blocks_count;
}

void Master :: markBlockAsDone(int block_id, int slave_id) {
    if(isValidBlock(block_id)) {
        in_progress.erase(block_id);
        current_info.blockToSlave[block_id] = slave_id;
        current_info.slaveToBlocks[slave_id] -> insert(block_id);
    }
}

int Master :: chooseNextBlockFor(int idle_slave) {
    SetIterator<int> iterator(prev_info.slaveToBlocks[idle_slave]);
    while(iterator.hasNext()) {
        int remoteBlock = iterator.next();
        if(todo.count(remoteBlock) == 1) {
            return remoteBlock;
        }
    }
    
    return *todo.begin();
}
/**************************************************************************************************************/


// Signals handlers
/**************************************************************************************************************/
void Master :: recieveReadyToWorkSignal(int& block_id, int& idle_slave) {
    MPI_Recv(&block_id, 1, MPI_INT, MPI_ANY_SOURCE, Signals :: READY_TO_WORK, MPI_COMM_WORLD, &status);
    idle_slave = status.MPI_SOURCE;
}

void Master :: sendBlockAssignmentSignal(int block_id, int idle_slave) {
    MPI_Send(&block_id, 1, MPI_INT, idle_slave, Signals :: ASSIGNED_BLOCK, MPI_COMM_WORLD); // may be better
}

void Master :: broadcastSignal(Signals :: CommunicationSignals signal) {
    for(int i = 1; i < processes_count; ++i) {
        MPI_Send(&nothing, 1, MPI_INT, i, signal, MPI_COMM_WORLD); 
    }
}

void Master :: sendNextIterationSignal() {
    broadcastSignal(Signals :: NEXT_ITERATION);
}

void Master :: sendTerminationSignal() {
    broadcastSignal(Signals :: TERMINATE);
}
/**************************************************************************************************************/


// Constructor
/**************************************************************************************************************/
Master :: Master() {
    MPI_Comm_size(MPI_COMM_WORLD, &processes_count);
    current_info.initialize(processes_count);
    prev_info.initialize(processes_count);
    block_info_server.setBlockInfoSource(&prev_info);
}
/**************************************************************************************************************/


// Serving block info
/**************************************************************************************************************/
void Master :: serveBlockInfo() {
    block_info_server.start();
}
/**************************************************************************************************************/


// IterationInfo methods
/**************************************************************************************************************/
void Master :: IterationInfo :: initialize(int processes_count) {
    for(int i = 0; i < processes_count; ++i) {
        slaveToBlocks[i] = new set<int>();
    }
}

void Master :: IterationInfo :: swapContent(IterationInfo& a, IterationInfo& b) {
    a.blockToSlave.swap(b.blockToSlave);
    a.slaveToBlocks.swap(b.slaveToBlocks);
}

void Master :: IterationInfo :: clear() {
    blockToSlave.clear();
    MapValuesIterator<int, set<int>*> iterator(&slaveToBlocks);
    while(iterator.hasNext()) {
        iterator.next() -> clear();
    }
}

Master :: IterationInfo :: ~IterationInfo() {
    MapValuesIterator<int, set<int>*> iterator(&slaveToBlocks);
    while(iterator.hasNext()) {
        delete iterator.next();
    }
}
/**************************************************************************************************************/


// BlockInfoServer methods
/**************************************************************************************************************/
void Master :: BlockInfoServer :: listenForRequest() {
    MPI_Irecv(&requested_block_id, 1, MPI_INT, MPI_ANY_SOURCE, Signals :: WHERE_IS_BLOCK, MPI_COMM_WORLD, &request);
}

void Master :: BlockInfoServer :: sendResponse() {
    int block_location = info_source -> blockToSlave[requested_block_id];
    MPI_Send(&block_location, 1, MPI_INT, status.MPI_SOURCE, Signals :: BLOCK_LOCATION, MPI_COMM_WORLD);
}
/**************************************************************************************************************/
